package com.blyat.xsoft.gateway.security.core;

import com.blyat.xsoft.gateway.security.refresher.JwtReactiveTokenRefresher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.security.authentication.AccountExpiredException;
import org.springframework.security.authentication.BadCredentialsException;
import org.springframework.security.authentication.ReactiveAuthenticationManager;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.oauth2.server.resource.BearerTokenAuthenticationToken;
import reactor.core.publisher.Mono;

/**
 * @Auther: syh
 * @Date: 2020/8/27
 * @Description:
 */
public class XSoftReactiveAuthenticationManager implements ReactiveAuthenticationManager {
    private static final Logger logger = LoggerFactory.getLogger(XSoftReactiveAuthenticationManager.class);

    private TokenStore tokenStore;
    // 自定义的token刷新处理器
    private JwtReactiveTokenRefresher executor;

    public XSoftReactiveAuthenticationManager(TokenStore tokenStore) {
        this.tokenStore = tokenStore;
    }

    public void setExecutor(JwtReactiveTokenRefresher executor) {
        this.executor = executor;
    }

    @Override
    public Mono<Authentication> authenticate(Authentication authentication) {
        return Mono.justOrEmpty(authentication)
                .filter(a -> a instanceof BearerTokenAuthenticationToken)
                .cast(BearerTokenAuthenticationToken.class)
                .map(BearerTokenAuthenticationToken::getToken)
                .flatMap(this::checkAccessToken)
                .cast(Authentication.class);
    }

    private Mono<OAuth2Authentication> checkAccessToken(final String accessToken) {
        OAuth2AccessToken oAuth2AccessToken = this.tokenStore.readAccessToken(accessToken);
        if (executor != null) {
            executor.setAccessToken(oAuth2AccessToken);
            if (executor.shouldRefresh()) {
                try {
                    logger.info("refresh token.");
                    return executor.refresh().map(token -> {
                        if (!token.equals(accessToken)) {
                            tokenStore.removeAccessToken(oAuth2AccessToken);
                        }
                        return token;
                    }).flatMap(token -> {
                        OAuth2Authentication oAuth2Authentication = this.tokenStore.readAuthentication(token);
                        if (oAuth2Authentication == null) {
                            return Mono.error(new BadCredentialsException("Token 无效!"));
                        } else {
                            return Mono.just(oAuth2Authentication);
                        }
                    });
                } catch (Exception e) {
                    logger.error("token refresh failed.", e);
                    return Mono.error(new AccountExpiredException("Token 失效!"));
                }
            }
        } else {
            if (oAuth2AccessToken == null) {
                return Mono.error(new BadCredentialsException("Token 无效!"));
            } else if (oAuth2AccessToken.isExpired()) {
                return Mono.error(new AccountExpiredException("Token 失效!"));
            }
        }

        OAuth2Authentication oAuth2Authentication = this.tokenStore.readAuthentication(accessToken);
        if (oAuth2Authentication == null) {
            return Mono.error(new BadCredentialsException("Token 无效!"));
        } else {
            return Mono.just(oAuth2Authentication);
        }
    }
}
